import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint

# 设置 Tushare token
ts.set_token('40a93ca838c790d9e0c1fdc4a804c6454383d5c436f71db1affa3a00')
pro = ts.pro_api()


# 获取股票数据（增加复权处理）
def get_stock_data(tickers, start_date, end_date, adj='qfq'):
    data_list = []
    for ticker in tickers:
        # 获取前复权数据
        df = pro.daily(ts_code=ticker, start_date=start_date, end_date=end_date)
        df['trade_date'] = pd.to_datetime(df['trade_date'])
        df.set_index('trade_date', inplace=True)
        df.sort_index(inplace=True)

        # 计算复权收盘价
        if adj == 'qfq':
            if 'adj_factor' in df.columns:
                df['close'] = df['close'] * df['adj_factor']
        df = df[['close']]
        df.columns = [ticker]
        data_list.append(df)
    return pd.concat(data_list, axis=1)


# 优化的动量策略：多周期动量融合
def momentum_strategy(data, short_window=5, medium_window=20, long_window=60):
    # 计算多周期动量
    short_momentum = data.pct_change(short_window)
    medium_momentum = data.pct_change(medium_window)
    long_momentum = data.pct_change(long_window)

    # 动量融合：给予不同周期不同权重
    combined_momentum = 0.4 * short_momentum + 0.3 * medium_momentum + 0.3 * long_momentum

    signals = pd.DataFrame(index=data.index, columns=data.columns)

    # 改进信号生成逻辑：加入动量强度过滤
    strong_momentum = combined_momentum > 0.01  # 正向强动量
    weak_momentum = (combined_momentum <= 0.01) & (combined_momentum >= -0.01)  # 弱动量
    strong_negative_momentum = combined_momentum < -0.01  # 负向强动量

    signals[strong_momentum] = 1  # 强正向动量时买入
    signals[weak_momentum] = 0  # 弱动量时持有或观望
    signals[strong_negative_momentum] = -1  # 强负向动量时卖出

    return signals


# 优化的LSTM策略：增强版价格走势预测
def lstm_strategy(data, ticker, look_back=60, epochs=20, batch_size=32):
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data[ticker].values.reshape(-1, 1))

    # 增加技术指标作为特征
    data['returns'] = data[ticker].pct_change()
    data['sma_20'] = data[ticker].rolling(window=20).mean()
    data['sma_50'] = data[ticker].rolling(window=50).mean()
    data['rsi_14'] = calculate_rsi(data[ticker], 14)
    data = data.dropna()

    # 准备多特征输入
    feature_data = data[[ticker, 'returns', 'sma_20', 'sma_50', 'rsi_14']].values
    scaled_features = scaler.fit_transform(feature_data)

    X, y = [], []
    for i in range(look_back, len(scaled_features)):
        X.append(scaled_features[i - look_back:i, :])
        y.append(scaled_features[i, 0])  # 预测目标为收盘价
    X, y = np.array(X), np.array(y)

    # 构建更复杂的LSTM模型
    model = Sequential()
    model.add(LSTM(units=50, return_sequences=True, input_shape=(X.shape[1], X.shape[2])))
    model.add(LSTM(units=50, return_sequences=False))
    model.add(Dense(units=25))
    model.add(Dense(units=1))
    model.compile(optimizer='adam', loss='mean_squared_error')

    # 改进训练过程：早停和模型检查点
    early_stop = EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True)
    checkpoint = ModelCheckpoint(f'best_{ticker}_model.h5', monitor='val_loss', save_best_only=True)

    model.fit(X, y, epochs=epochs, batch_size=batch_size, validation_split=0.2,
              callbacks=[early_stop, checkpoint], verbose=0)

    # 准备预测数据
    last_60 = data[[ticker, 'returns', 'sma_20', 'sma_50', 'rsi_14']].iloc[-look_back:].values
    last_60_scaled = scaler.transform(last_60)
    last_60_scaled = last_60_scaled.reshape(1, look_back, X.shape[2])

    # 预测下一时刻价格
    predicted_scaled = model.predict(last_60_scaled)
    predicted_price = scaler.inverse_transform(np.hstack((predicted_scaled, np.zeros((1, 4)))))[:, 0]

    last_price = data[ticker].iloc[-1]

    # 增加预测可信度判断
    confidence = calculate_prediction_confidence(model, X, y)
    if predicted_price > last_price and confidence > 0.6:
        return 1
    elif predicted_price < last_price and confidence > 0.6:
        return -1
    else:
        return 0  # 信心不足时不交易


# 计算RSI指标
def calculate_rsi(series, period=14):
    delta = series.diff()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)

    avg_gain = gain.rolling(window=period).mean()
    avg_loss = loss.rolling(window=period).mean()

    rs = avg_gain / avg_loss
    rsi = 100 - (100 / (1 + rs))
    return rsi


# 计算预测可信度
def calculate_prediction_confidence(model, X, y):
    # 简单方法：基于模型在验证集上的表现
    val_predictions = model.predict(X[-20:])
    mse = np.mean((val_predictions.flatten() - y[-20:]) ** 2)
    # 转换为0-1之间的可信度
    max_mse = np.var(y)
    confidence = 1 - min(mse / max_mse, 1.0)
    return confidence


# 增强的量化交易系统类
class TradingSystem:
    def __init__(self, data):
        self.data = data
        self.portfolio = pd.DataFrame(index=data.index, columns=['Value'])
        self.portfolio['Value'] = 1.0
        self.stop_loss = -0.08  # 优化止损比例
        self.take_profit = 0.15  # 优化止盈比例
        self.max_drawdown = 0.15  # 最大回撤限制

    def generate_signals(self, strategy, **kwargs):
        signals = pd.DataFrame(index=self.data.index, columns=self.data.columns)
        if strategy == 'momentum':
            signals = momentum_strategy(self.data, **kwargs)
        elif strategy == 'lstm':
            signals = pd.DataFrame(index=self.data.index, columns=self.data.columns)
            for ticker in self.data.columns:
                signal = lstm_strategy(self.data, ticker, **kwargs)
                signals[ticker] = signal
        return signals

    def backtest(self, strategy, **kwargs):
        signals = self.generate_signals(strategy, **kwargs)
        returns = self.data.pct_change()

        # 考虑交易成本
        transaction_cost = 0.001  # 千分之一的交易成本
        strategy_returns = (signals.shift(1).fillna(0) * returns).sum(axis=1)
        # 计算交易信号变化带来的成本
        position_changes = signals.diff().fillna(signals).abs().sum(axis=1)
        strategy_returns -= position_changes * transaction_cost

        portfolio = pd.DataFrame(index=self.data.index, columns=['Value'])
        portfolio['Value'] = 1.0

        # 记录最大回撤
        max_value = 1.0

        for i in range(1, len(portfolio)):
            current_value = portfolio['Value'].iloc[i - 1]
            daily_return = strategy_returns.iloc[i]

            # 计算当日价值
            new_value = current_value * (1 + daily_return)

            # 最大回撤控制
            max_value = max(max_value, current_value)
            drawdown = (current_value - max_value) / max_value
            if drawdown < -self.max_drawdown:
                # 超过最大回撤时清仓
                new_value = current_value * (1 - daily_return)  # 避免进一步损失

            # 止损止盈
            if new_value < current_value * (1 + self.stop_loss):
                new_value = current_value * (1 + self.stop_loss)
            elif new_value > current_value * (1 + self.take_profit):
                new_value = current_value * (1 + self.take_profit)

            portfolio.loc[self.portfolio.index[i], 'Value'] = new_value

        return portfolio

    def plot_results(self, portfolios):
        plt.figure(figsize=(12, 7))
        for strategy, portfolio in portfolios.items():
            plt.plot(portfolio['Value'], label=strategy)

        # 计算并显示关键指标
        plt.axhline(y=1.0, color='r', linestyle='--', alpha=0.3)

        plt.title('Strategy Performance Comparison')
        plt.xlabel('Date')
        plt.ylabel('Portfolio Value')
        plt.legend()

        # 添加指标文本框
        metrics_text = ""
        for strategy, portfolio in portfolios.items():
            metrics = self.calculate_metrics(portfolio)
            metrics_text += f"{strategy}: 总收益 {metrics['总收益率']}, 年化 {metrics['年化收益率']}, 最大回撤 {metrics['最大回撤']}, 夏普比率 {metrics['夏普比率']}\n"

        plt.figtext(0.01, 0.01, metrics_text, fontsize=9,
                    bbox=dict(facecolor='white', alpha=0.8))

        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()

    def calculate_metrics(self, portfolio):
        """计算策略的关键绩效指标"""
        returns = portfolio['Value'].pct_change().dropna()
        total_return = (portfolio['Value'].iloc[-1] - 1) * 100
        annual_return = ((1 + total_return / 100) ** (252 / len(returns)) - 1) * 100
        max_drawdown = (1 - portfolio['Value'] / portfolio['Value'].cummax()).max() * 100
        sharpe = (annual_return / 252 ** 0.5) / (returns.std() * 252 ** 0.5) if returns.std() != 0 else 0

        metrics = {
            '总收益率': f"{total_return:.2f}%",
            '年化收益率': f"{annual_return:.2f}%",
            '最大回撤': f"{max_drawdown:.2f}%",
            '夏普比率': f"{sharpe:.2f}"
        }
        return metrics

    def compare_strategies(self, portfolio1, name1, portfolio2, name2):
        """比较两个策略的绩效并返回更优策略"""
        metrics1 = self.calculate_metrics(portfolio1)
        metrics2 = self.calculate_metrics(portfolio2)

        # 转换指标为数值进行比较
        def parse_metric(metric_str):
            return float(metric_str.strip('%'))

        # 比较总收益率
        total_return1 = parse_metric(metrics1['总收益率'])
        total_return2 = parse_metric(metrics2['总收益率'])

        # 比较夏普比率
        sharpe1 = parse_metric(metrics1['夏普比率'])
        sharpe2 = parse_metric(metrics2['夏普比率'])

        # 比较最大回撤(负向指标，越小越好)
        drawdown1 = parse_metric(metrics1['最大回撤'])
        drawdown2 = parse_metric(metrics2['最大回撤'])

        # 计算综合评分 (权重可根据投资偏好调整)
        score1 = 0.5 * total_return1 + 0.3 * sharpe1 - 0.2 * drawdown1
        score2 = 0.5 * total_return2 + 0.3 * sharpe2 - 0.2 * drawdown2

        # 输出比较结果
        print("\n策略对比分析:")
        print(f"{name1} vs {name2}")
        print(
            f"{name1}: 总收益 {metrics1['总收益率']}, 夏普比率 {metrics1['夏普比率']}, 最大回撤 {metrics1['最大回撤']}")
        print(
            f"{name2}: 总收益 {metrics2['总收益率']}, 夏普比率 {metrics2['夏普比率']}, 最大回撤 {metrics2['最大回撤']}")

        if score1 > score2:
            print(f"\n结论: {name1} 策略更优，综合评分: {score1:.2f} vs {score2:.2f}")
            return name1, portfolio1, score1
        elif score2 > score1:
            print(f"\n结论: {name2} 策略更优，综合评分: {score2:.2f} vs {score1:.2f}")
            return name2, portfolio2, score2
        else:
            print(f"\n结论: 两个策略表现相当，综合评分: {score1:.2f}")
            return None, None, 0


if __name__ == "__main__":
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

    tickers_list = ['000001.SZ', '600000.SH', '601318.SH', '000002.SZ']
    start_date = '20190101'  # 调整回测起始时间
    end_date = '20231231'  # 调整回测结束时间

    data = get_stock_data(tickers_list, start_date, end_date)

    ts = TradingSystem(data)

    # 回测优化后的动量策略
    momentum_portfolio = ts.backtest('momentum', short_window=5, medium_window=20, long_window=60)
    momentum_metrics = ts.calculate_metrics(momentum_portfolio)
    print("动量策略绩效指标:")
    for metric, value in momentum_metrics.items():
        print(f"{metric}: {value}")

    # 回测优化后的LSTM策略
    lstm_portfolio = ts.backtest('lstm', look_back=60, epochs=30, batch_size=32)
    lstm_metrics = ts.calculate_metrics(lstm_portfolio)
    print("\nLSTM策略绩效指标:")
    for metric, value in lstm_metrics.items():
        print(f"{metric}: {value}")

    # 比较两个策略
    best_strategy, best_portfolio, best_score = ts.compare_strategies(
        momentum_portfolio, "动量策略",
        lstm_portfolio, "LSTM策略"
    )

    # 绘制策略对比结果
    portfolios = {'动量策略': momentum_portfolio, 'LSTM策略': lstm_portfolio}
    ts.plot_results(portfolios)